package com.yc.cloud.openai.portal.service.impl;

import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Map;
import java.util.Objects;

import org.springframework.stereotype.Service;

import com.yc.cloud.openai.portal.config.ComfyUIProperties;
import com.yc.cloud.openai.portal.dao.impl.ComfyUiClientImpl;
import com.yc.cloud.openai.portal.dto.request.TextToImageRequest;
import com.yc.cloud.openai.portal.dto.request.TextToVideoRequest;
import com.yc.cloud.openai.portal.service.WorkflowProcessService;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import lombok.SneakyThrows;
import lombok.val;
import lombok.extern.slf4j.Slf4j;

@Service
@Slf4j
public class WorkflowProcessServiceImpl implements WorkflowProcessService {

    private final ComfyUIProperties comfyUIProperties;

    public WorkflowProcessServiceImpl(ComfyUIProperties comfyUIProperties) {
        this.comfyUIProperties = comfyUIProperties;
    }

    @Override
    @SneakyThrows
    @SuppressWarnings("unchecked")
    public <T> String processWorkflow(String workflow, T request) {
        Map<String, Object> workflowData = JSONUtil.toBean(workflow, Map.class);
        if (request instanceof TextToImageRequest) {

            for (Map.Entry<String, Object> entry : workflowData.entrySet()) {
                Map<String, Object> node = (Map<String, Object>) entry.getValue();
                processTextToImageNode(node, (TextToImageRequest) request);
            }
            } else if (request instanceof TextToVideoRequest) {
                for (Map.Entry<String, Object> entry : workflowData.entrySet()) {
                    Map<String, Object> node = (Map<String, Object>) entry.getValue();
                    processTextToVideoNode(node, (TextToVideoRequest) request);
                }
            } 

        return JSONUtil.toJsonPrettyStr(workflowData);
    }

    /**
     * 处理文本到图片的节点
     * 
     * @param node
     * @param request
     */
    @SuppressWarnings("unchecked")
    private void processTextToImageNode(Map<String, Object> node, TextToImageRequest request) {
        String classType = (String) node.get("class_type");
        Map<String, Object> inputs = (Map<String, Object>) node.get("inputs");
        if (inputs == null)
            return;

        // 处理文本节点
        if ((classType.equals("CLIPTextEncode") || classType.equals("DeepTranslatorCLIPTextEncodeNode"))
                && inputs.containsKey("text")) {
            String oldText = (String) inputs.get("text");
            if (!oldText.contains("watermark")) {
                String prompt = ((TextToImageRequest) request).getPrompt();
                inputs.put("text", prompt);
                log.info("原提示词：{}, 新提示词：{}", oldText, prompt);
            }
        }
        // 处理图片尺寸
        else if (classType.equals("EmptyLatentImage")) {
            if (request instanceof TextToImageRequest) {
                TextToImageRequest imageRequest = (TextToImageRequest) request;
                if (inputs.containsKey("width")) {
                    val oldWidth = inputs.get("width");
                    inputs.put("width", imageRequest.getWidth());
                    log.info("原宽度：{}, 新宽度：{}", oldWidth, imageRequest.getWidth());
                }
                if (inputs.containsKey("height")) {
                    val oldHeight = inputs.get("height");
                    inputs.put("height", imageRequest.getHeight());
                    log.info("原高度：{}, 新高度：{}", oldHeight, imageRequest.getHeight());
                }
            }
        }

        // 处理随机种子
        else if (classType.equals("RandomNoise") && inputs.containsKey("noise_seed")) {
            inputs.put("noise_seed", DateUtil.current());
        }

        // 处理其他种子
        else if (inputs.containsKey("seed")) {
            Object oldSeed = inputs.get("seed");
            inputs.put("seed", DateUtil.current());
            log.info("原 seed：{}, 新 seed：{}", oldSeed, inputs.get("seed"));
        }

        // 处理 Lora
        else if (inputs.containsKey("lora_name")) {
            if (request instanceof TextToImageRequest) {
                val oldLoraName = inputs.get("lora_name");
                var newLoraName = ((TextToImageRequest) request).getStyle().getValue();
                if (StrUtil.isNotEmpty(newLoraName)) {
                    newLoraName = comfyUIProperties.getPrefixLora() + newLoraName;
                } else {
                    newLoraName = "None";
                }
                inputs.put("lora_name", newLoraName);
                log.info("原 lora_name：{}, 新 lora_name：{}", oldLoraName, newLoraName);
            }
        }
    }

    @SuppressWarnings("unchecked")
    private void processTextToVideoNode(Map<String, Object> node, TextToVideoRequest request) {
        String classType = (String) node.get("class_type");
        Map<String, Object> inputs = (Map<String, Object>) node.get("inputs");
        if (inputs == null) return;

        // 处理提示词
        if (classType.equals("CLIPTextEncode") && inputs.containsKey("text")) {
            String oldText = (String) inputs.get("text");
            inputs.put("text", request.getPrompt());
            log.info("原提示词：{}, 新提示词：{}", oldText, request.getPrompt());
        }
        
        // 处理随机种子
        else if (classType.equals("RandomNoise") && inputs.containsKey("noise_seed")) {
            Object oldSeed = inputs.get("noise_seed");
            inputs.put("noise_seed", DateUtil.current());
            log.info("原 noise_seed：{}, 新 noise_seed：{}", oldSeed, inputs.get("noise_seed"));
        }
    }

    /**
     * 读取工作流信息
     *
     * @param workflowName 工作流文件名
     * @return 工作流信息
     */
    @SneakyThrows
    public String readWorkflow(String workflowName) {
        Path jsonPath;
        jsonPath = Paths.get(Objects
                .requireNonNull(
                        ComfyUiClientImpl.class.getClassLoader().getResource("comfyui/workflows/" + workflowName))
                .toURI());
        // 读取文件内容
        return new String(Files.readAllBytes(jsonPath));
    }

}
